package com.example.demo.controller;

import jakarta.servlet.ServletOutputStream;
import jakarta.servlet.http.HttpServletResponse;
import okhttp3.*;
import org.json.JSONArray;
import org.json.JSONObject;
import org.springframework.web.bind.annotation.CrossOrigin;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RestController;

import java.io.IOException;
import java.io.InputStream;

@RestController
@CrossOrigin
public class DemoController {

    private static final String API_KEY = "";
    private static final String SECRET_KEY = "";
    private final OkHttpClient HTTP_CLIENT = new OkHttpClient().newBuilder().build();

    @PostMapping("/test")
    public void demo(@RequestBody String message , HttpServletResponse res) throws IOException {
        MediaType mediaType = MediaType.parse("application/json");
        JSONObject jsonObject = new JSONObject();
        jsonObject.put("messages", new JSONArray(message));
        jsonObject.put("stream",true);
        okhttp3.RequestBody body = okhttp3.RequestBody.create(mediaType, jsonObject.toString());
        Request request = new Request.Builder()
                .url("https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie-speed-128k?access_token=" + getAccessToken())
                .method("POST", body)
                .addHeader("Content-Type", "application/json")
                .build();
        Response response = HTTP_CLIENT.newCall(request).execute();
//        System.out.println(response.body().string());
        //流式返回
        res.addHeader("Content-type","text/event-stream; charset=utf-8");
        res.addHeader("Cache-Control", "no-cache");
        ServletOutputStream outputStream = res.getOutputStream();
        if(response.isSuccessful()) {
            okhttp3.ResponseBody responseBody = response.body();
            if(responseBody != null) {
                InputStream inputStream = responseBody.byteStream();
                byte[] buffer = new byte[1024];
                int bytesRead = 0;
                while((bytesRead=inputStream.read(buffer))!=-1){
                    String data = new String(buffer, 0, bytesRead);
                    String[] lines = data.split("\n");
                    for(String line : lines){
                        if(line.startsWith("data: ")){
                            //验证是否是流式返回
                            JSONObject jsonObject1 = new JSONObject(data.substring(6));
                            System.out.print(jsonObject1.get("result"));
                            outputStream.write((data.substring(6)/*+'\n'*/).getBytes());
                            outputStream.flush();
                        }
                    }
                }
            }

        }
        outputStream.close();

    }

    private String getAccessToken() throws IOException {
        MediaType mediaType = MediaType.parse("application/x-www-form-urlencoded");
        okhttp3.RequestBody body = okhttp3.RequestBody.create(mediaType, "grant_type=client_credentials&client_id=" + API_KEY
                + "&client_secret=" + SECRET_KEY);
        Request request = new Request.Builder()
                .url("https://aip.baidubce.com/oauth/2.0/token")
                .method("POST", body)
                .addHeader("Content-Type", "application/x-www-form-urlencoded")
                .build();
        Response response = HTTP_CLIENT.newCall(request).execute();
        return new JSONObject(response.body().string()).getString("access_token");
    }
}
